import numpy as np
import copy



def transform(a,action_space=None,index=False):
    """
    如果index为False,则输出实际动作
    否则，输出动作对应序号
    """
    if index:
        return np.where(np.array(action_space)==a)[0][0]

    return  action_space[int(a)]

def select_A(s,action_space=None,pie=None,b=None,e=0.1):
    """
    输入现状态s，pie，b
    对策略pie实行贪婪原则,pie:{s:np.array([a1,a2,....])}
    对策略b实行ε-贪婪原则,b:{s:np.array([a1,a2,....])}
    来实现同轨、离轨策略的训练

    如果选择离轨策略，即根据b选择动作，择参数e失效（利用argmax函数决策）
    返回 动作序号->动作
    """
    
    if pie:
        alternative_a=pie[s]
        ass=np.where(alternative_a==np.max(alternative_a))[0]
        
        a=np.random.choice(ass)

        p0=np.random.rand()
        if p0<1-e:
            act_a=transform(a,action_space)
            return act_a
        else:
            ranges=np.arange(len(action_space))
            idxes=list(set(ranges)-set(ass))
            if idxes:
                a=np.random.choice(idxes)
                act_a=transform(a,action_space)
                return act_a
            else:
                act_a=transform(a,action_space)
                return act_a

    if b:
        alternative_a=b[s]
        a=np.argmax(alternative_a)
        p0=np.random.rand()
        if p0<e:
            ranges=list(range(len(alternative_a)))
            ranges.remove(a)

            a=np.random.choice(ranges)

            return transform(a,action_space)

        else:
            return transform(a,action_space)

def softmax(x):
    y=np.exp(x-x.max())

    return y/y.sum()


def update_stategy(pie=None,b=None,Q=None):
    """
    输入pie或者b,Q:{s:np.array([])}
    采取策略迭代和值迭代两种方式
    返回策略
    """
    if Q:
        if pie:
            for s in pie:
                alternative_act=np.array(pie[s])
                indexs=list(np.where(Q[s]==np.max(Q[s]))[0])
                alternative_act[indexs]=1
                indexs=list(set(list(range(alternative_act.shape[0])))-set(indexs))
                alternative_act[indexs]=0
                pie[s]=alternative_act
            return pie
        if b:
            for s in b:
                every_p=softmax(Q[s])
                b[s]=every_p

            return b
    else:
        raise ValueError





        